core/iter/adapters/
step_by.rs

1use crate::intrinsics;
2use crate::iter::{FusedIterator, TrustedLen, TrustedRandomAccess, from_fn};
3use crate::num::NonZero;
4use crate::ops::{Range, Try};
5use crate::range::RangeIter;
6
7/// An iterator for stepping iterators by a custom amount.
8///
9/// This `struct` is created by the [`step_by`] method on [`Iterator`]. See
10/// its documentation for more.
11///
12/// [`step_by`]: Iterator::step_by
13/// [`Iterator`]: trait.Iterator.html
14#[must_use = "iterators are lazy and do nothing unless consumed"]
15#[stable(feature = "iterator_step_by", since = "1.28.0")]
16#[derive(Clone, Debug)]
17pub struct StepBy<I> {
18    /// This field is guaranteed to be preprocessed by the specialized `SpecRangeSetup::setup`
19    /// in the constructor.
20    /// For most iterators that processing is a no-op, but for Range<{integer}> types it is lossy
21    /// which means the inner iterator cannot be returned to user code.
22    /// Additionally this type-dependent preprocessing means specialized implementations
23    /// cannot be used interchangeably.
24    iter: I,
25    /// This field is `step - 1`, aka the correct amount to pass to `nth` when iterating.
26    /// It MUST NOT be `usize::MAX`, as `unsafe` code depends on being able to add one
27    /// without the risk of overflow.  (This is important so that length calculations
28    /// don't need to check for division-by-zero, for example.)
29    step_minus_one: usize,
30    first_take: bool,
31}
32
33impl<I> StepBy<I> {
34    #[inline]
35    pub(in crate::iter) fn new(iter: I, step: usize) -> StepBy<I> {
36        assert!(step != 0);
37        let iter = <I as SpecRangeSetup<I>>::setup(iter, step);
38        StepBy { iter, step_minus_one: step - 1, first_take: true }
39    }
40
41    /// The `step` that was originally passed to `Iterator::step_by(step)`,
42    /// aka `self.step_minus_one + 1`.
43    #[inline]
44    fn original_step(&self) -> NonZero<usize> {
45        // SAFETY: By type invariant, `step_minus_one` cannot be `MAX`, which
46        // means the addition cannot overflow and the result cannot be zero.
47        unsafe { NonZero::new_unchecked(intrinsics::unchecked_add(self.step_minus_one, 1)) }
48    }
49}
50
51#[stable(feature = "iterator_step_by", since = "1.28.0")]
52impl<I> Iterator for StepBy<I>
53where
54    I: Iterator,
55{
56    type Item = I::Item;
57
58    #[inline]
59    fn next(&mut self) -> Option<Self::Item> {
60        self.spec_next()
61    }
62
63    #[inline]
64    fn size_hint(&self) -> (usize, Option<usize>) {
65        self.spec_size_hint()
66    }
67
68    #[inline]
69    fn nth(&mut self, n: usize) -> Option<Self::Item> {
70        self.spec_nth(n)
71    }
72
73    fn try_fold<Acc, F, R>(&mut self, acc: Acc, f: F) -> R
74    where
75        F: FnMut(Acc, Self::Item) -> R,
76        R: Try<Output = Acc>,
77    {
78        self.spec_try_fold(acc, f)
79    }
80
81    #[inline]
82    fn fold<Acc, F>(self, acc: Acc, f: F) -> Acc
83    where
84        F: FnMut(Acc, Self::Item) -> Acc,
85    {
86        self.spec_fold(acc, f)
87    }
88}
89
90impl<I> StepBy<I>
91where
92    I: ExactSizeIterator,
93{
94    // The zero-based index starting from the end of the iterator of the
95    // last element. Used in the `DoubleEndedIterator` implementation.
96    fn next_back_index(&self) -> usize {
97        let rem = self.iter.len() % self.original_step();
98        if self.first_take { if rem == 0 { self.step_minus_one } else { rem - 1 } } else { rem }
99    }
100}
101
102#[stable(feature = "double_ended_step_by_iterator", since = "1.38.0")]
103impl<I> DoubleEndedIterator for StepBy<I>
104where
105    I: DoubleEndedIterator + ExactSizeIterator,
106{
107    #[inline]
108    fn next_back(&mut self) -> Option<Self::Item> {
109        self.spec_next_back()
110    }
111
112    #[inline]
113    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
114        self.spec_nth_back(n)
115    }
116
117    fn try_rfold<Acc, F, R>(&mut self, init: Acc, f: F) -> R
118    where
119        F: FnMut(Acc, Self::Item) -> R,
120        R: Try<Output = Acc>,
121    {
122        self.spec_try_rfold(init, f)
123    }
124
125    #[inline]
126    fn rfold<Acc, F>(self, init: Acc, f: F) -> Acc
127    where
128        Self: Sized,
129        F: FnMut(Acc, Self::Item) -> Acc,
130    {
131        self.spec_rfold(init, f)
132    }
133}
134
135// StepBy can only make the iterator shorter, so the len will still fit.
136#[stable(feature = "iterator_step_by", since = "1.28.0")]
137impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {}
138
139// StepBy stops yielding items once the underlying iterator does, so it is fused
140// whenever the underlying iterator is fused.
141#[stable(feature = "step_by_fused", since = "1.99.0")]
142impl<I> FusedIterator for StepBy<I> where I: FusedIterator {}
143
144// SAFETY: This adapter is shortening. TrustedLen requires the upper bound to be calculated correctly.
145// These requirements can only be satisfied when the upper bound of the inner iterator's upper
146// bound is never `None`. I: TrustedRandomAccess happens to provide this guarantee while
147// I: TrustedLen would not.
148// This also covers the Range specializations since the ranges also implement TRA
149#[unstable(feature = "trusted_len", issue = "37572")]
150unsafe impl<I> TrustedLen for StepBy<I> where I: Iterator + TrustedRandomAccess {}
151
152trait SpecRangeSetup<T> {
153    fn setup(inner: T, step: usize) -> T;
154}
155
156impl<T> SpecRangeSetup<T> for T {
157    #[inline]
158    default fn setup(inner: T, _step: usize) -> T {
159        inner
160    }
161}
162
163/// Specialization trait to optimize `StepBy<Range<{integer}>>` iteration.
164///
165/// # Safety
166///
167/// Technically this is safe to implement (look ma, no unsafe!), but in reality
168/// a lot of unsafe code relies on ranges over integers being correct.
169///
170/// For correctness *all* public StepBy methods must be specialized
171/// because `setup` drastically alters the meaning of the struct fields so that mixing
172/// different implementations would lead to incorrect results.
173unsafe trait StepByImpl<I> {
174    type Item;
175
176    fn spec_next(&mut self) -> Option<Self::Item>;
177
178    fn spec_size_hint(&self) -> (usize, Option<usize>);
179
180    fn spec_nth(&mut self, n: usize) -> Option<Self::Item>;
181
182    fn spec_try_fold<Acc, F, R>(&mut self, acc: Acc, f: F) -> R
183    where
184        F: FnMut(Acc, Self::Item) -> R,
185        R: Try<Output = Acc>;
186
187    fn spec_fold<Acc, F>(self, acc: Acc, f: F) -> Acc
188    where
189        F: FnMut(Acc, Self::Item) -> Acc;
190}
191
192/// Specialization trait for double-ended iteration.
193///
194/// See also: `StepByImpl`
195///
196/// # Safety
197///
198/// The specializations must be implemented together with `StepByImpl`
199/// where applicable. I.e. if `StepBy` does support backwards iteration
200/// for a given iterator and that is specialized for forward iteration then
201/// it must also be specialized for backwards iteration.
202unsafe trait StepByBackImpl<I> {
203    type Item;
204
205    fn spec_next_back(&mut self) -> Option<Self::Item>
206    where
207        I: DoubleEndedIterator + ExactSizeIterator;
208
209    fn spec_nth_back(&mut self, n: usize) -> Option<Self::Item>
210    where
211        I: DoubleEndedIterator + ExactSizeIterator;
212
213    fn spec_try_rfold<Acc, F, R>(&mut self, init: Acc, f: F) -> R
214    where
215        I: DoubleEndedIterator + ExactSizeIterator,
216        F: FnMut(Acc, Self::Item) -> R,
217        R: Try<Output = Acc>;
218
219    fn spec_rfold<Acc, F>(self, init: Acc, f: F) -> Acc
220    where
221        I: DoubleEndedIterator + ExactSizeIterator,
222        F: FnMut(Acc, Self::Item) -> Acc;
223}
224
225unsafe impl<I: Iterator> StepByImpl<I> for StepBy<I> {
226    type Item = I::Item;
227
228    #[inline]
229    default fn spec_next(&mut self) -> Option<I::Item> {
230        let step_size = if self.first_take { 0 } else { self.step_minus_one };
231        self.first_take = false;
232        self.iter.nth(step_size)
233    }
234
235    #[inline]
236    default fn spec_size_hint(&self) -> (usize, Option<usize>) {
237        #[inline]
238        fn first_size(step: NonZero<usize>) -> impl Fn(usize) -> usize {
239            move |n| if n == 0 { 0 } else { 1 + (n - 1) / step }
240        }
241
242        #[inline]
243        fn other_size(step: NonZero<usize>) -> impl Fn(usize) -> usize {
244            move |n| n / step
245        }
246
247        let (low, high) = self.iter.size_hint();
248
249        if self.first_take {
250            let f = first_size(self.original_step());
251            (f(low), high.map(f))
252        } else {
253            let f = other_size(self.original_step());
254            (f(low), high.map(f))
255        }
256    }
257
258    #[inline]
259    default fn spec_nth(&mut self, mut n: usize) -> Option<I::Item> {
260        if self.first_take {
261            self.first_take = false;
262            let first = self.iter.next()?;
263            if n == 0 {
264                return Some(first);
265            }
266            n -= 1;
267        }
268        // n and self.step_minus_one are indices, we need to add 1 to get the amount of elements
269        // When calling `.nth`, we need to subtract 1 again to convert back to an index
270        let mut step = self.original_step().get();
271        // n + 1 could overflow
272        // thus, if n is usize::MAX, instead of adding one, we call .nth(step)
273        if n == usize::MAX {
274            self.iter.nth(step - 1)?;
275        } else {
276            n += 1;
277        }
278
279        // overflow handling
280        loop {
281            let mul = n.checked_mul(step);
282            {
283                if intrinsics::likely(mul.is_some()) {
284                    return self.iter.nth(mul.unwrap() - 1);
285                }
286            }
287            let div_n = usize::MAX / n;
288            let div_step = usize::MAX / step;
289            let nth_n = div_n * n;
290            let nth_step = div_step * step;
291            let nth = if nth_n > nth_step {
292                step -= div_n;
293                nth_n
294            } else {
295                n -= div_step;
296                nth_step
297            };
298
299            self.iter.nth(nth - 1)?;
300        }
301    }
302
303    default fn spec_try_fold<Acc, F, R>(&mut self, mut acc: Acc, mut f: F) -> R
304    where
305        F: FnMut(Acc, Self::Item) -> R,
306        R: Try<Output = Acc>,
307    {
308        #[inline]
309        fn nth<I: Iterator>(
310            iter: &mut I,
311            step_minus_one: usize,
312        ) -> impl FnMut() -> Option<I::Item> + '_ {
313            move || iter.nth(step_minus_one)
314        }
315
316        if self.first_take {
317            self.first_take = false;
318            match self.iter.next() {
319                None => return try { acc },
320                Some(x) => acc = f(acc, x)?,
321            }
322        }
323        from_fn(nth(&mut self.iter, self.step_minus_one)).try_fold(acc, f)
324    }
325
326    default fn spec_fold<Acc, F>(mut self, mut acc: Acc, mut f: F) -> Acc
327    where
328        F: FnMut(Acc, Self::Item) -> Acc,
329    {
330        #[inline]
331        fn nth<I: Iterator>(
332            iter: &mut I,
333            step_minus_one: usize,
334        ) -> impl FnMut() -> Option<I::Item> + '_ {
335            move || iter.nth(step_minus_one)
336        }
337
338        if self.first_take {
339            self.first_take = false;
340            match self.iter.next() {
341                None => return acc,
342                Some(x) => acc = f(acc, x),
343            }
344        }
345        from_fn(nth(&mut self.iter, self.step_minus_one)).fold(acc, f)
346    }
347}
348
349unsafe impl<I: DoubleEndedIterator + ExactSizeIterator> StepByBackImpl<I> for StepBy<I> {
350    type Item = I::Item;
351
352    #[inline]
353    default fn spec_next_back(&mut self) -> Option<Self::Item> {
354        self.iter.nth_back(self.next_back_index())
355    }
356
357    #[inline]
358    default fn spec_nth_back(&mut self, n: usize) -> Option<I::Item> {
359        // `self.iter.nth_back(usize::MAX)` does the right thing here when `n`
360        // is out of bounds because the length of `self.iter` does not exceed
361        // `usize::MAX` (because `I: ExactSizeIterator`) and `nth_back` is
362        // zero-indexed
363        let n = n.saturating_mul(self.original_step().get()).saturating_add(self.next_back_index());
364        self.iter.nth_back(n)
365    }
366
367    default fn spec_try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R
368    where
369        F: FnMut(Acc, Self::Item) -> R,
370        R: Try<Output = Acc>,
371    {
372        #[inline]
373        fn nth_back<I: DoubleEndedIterator>(
374            iter: &mut I,
375            step_minus_one: usize,
376        ) -> impl FnMut() -> Option<I::Item> + '_ {
377            move || iter.nth_back(step_minus_one)
378        }
379
380        match self.next_back() {
381            None => try { init },
382            Some(x) => {
383                let acc = f(init, x)?;
384                from_fn(nth_back(&mut self.iter, self.step_minus_one)).try_fold(acc, f)
385            }
386        }
387    }
388
389    #[inline]
390    default fn spec_rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc
391    where
392        Self: Sized,
393        F: FnMut(Acc, I::Item) -> Acc,
394    {
395        #[inline]
396        fn nth_back<I: DoubleEndedIterator>(
397            iter: &mut I,
398            step_minus_one: usize,
399        ) -> impl FnMut() -> Option<I::Item> + '_ {
400            move || iter.nth_back(step_minus_one)
401        }
402
403        match self.next_back() {
404            None => init,
405            Some(x) => {
406                let acc = f(init, x);
407                from_fn(nth_back(&mut self.iter, self.step_minus_one)).fold(acc, f)
408            }
409        }
410    }
411}
412
413/// For these implementations, `SpecRangeSetup` calculates the number
414/// of iterations that will be needed and stores that in `iter.end`.
415///
416/// The various iterator implementations then rely on that to not need
417/// overflow checking, letting loops just be counted instead.
418///
419/// These only work for unsigned types, and will need to be reworked
420/// if you want to use it to specialize on signed types.
421///
422/// Currently these are only implemented for integers up to `usize` due to
423/// correctness issues around `ExactSizeIterator` impls on 16bit platforms.
424/// And since `ExactSizeIterator` is a prerequisite for backwards iteration
425/// and we must consistently specialize backwards and forwards iteration
426/// that makes the situation complicated enough that it's not covered
427/// for now.
428///
429/// After `SpecRangeSetup::setup`, both `Range<T>` and its new-range wrapper
430/// `RangeIter<T>` carry the cursor and countdown in the same underlying legacy
431/// `Range`. This accessor exposes that shared range so one specialization can
432/// serve both: it is an identity for `Range<T>` and unwraps the newtype for
433/// `RangeIter<T>`, so it compiles away.
434trait AsLegacyRange<T> {
435    fn as_legacy_range(&self) -> &Range<T>;
436    fn as_legacy_range_mut(&mut self) -> &mut Range<T>;
437}
438
439impl<T> AsLegacyRange<T> for Range<T> {
440    #[inline]
441    fn as_legacy_range(&self) -> &Range<T> {
442        self
443    }
444    #[inline]
445    fn as_legacy_range_mut(&mut self) -> &mut Range<T> {
446        self
447    }
448}
449
450impl<T> AsLegacyRange<T> for RangeIter<T> {
451    #[inline]
452    fn as_legacy_range(&self) -> &Range<T> {
453        &self.0
454    }
455    #[inline]
456    fn as_legacy_range_mut(&mut self) -> &mut Range<T> {
457        &mut self.0
458    }
459}
460
461macro_rules! spec_int_ranges {
462    ($ctor:ident; $($t:ty)*) => ($(
463
464        const _: () = assert!(usize::BITS >= <$t>::BITS);
465
466        impl SpecRangeSetup<$ctor<$t>> for $ctor<$t> {
467            #[inline]
468            fn setup(mut r: $ctor<$t>, step: usize) -> $ctor<$t> {
469                let inner_len = r.size_hint().0;
470                // If step exceeds $t::MAX, then the count will be at most 1 and
471                // thus always fit into $t.
472                let yield_count = inner_len.div_ceil(step);
473                // Turn the range end into an iteration counter
474                r.as_legacy_range_mut().end = yield_count as $t;
475                r
476            }
477        }
478
479        unsafe impl StepByImpl<$ctor<$t>> for StepBy<$ctor<$t>> {
480            #[inline]
481            fn spec_next(&mut self) -> Option<$t> {
482                // if a step size larger than the type has been specified fall back to
483                // t::MAX, in which case remaining will be at most 1.
484                let step = <$t>::try_from(self.original_step().get()).unwrap_or(<$t>::MAX);
485                let r = self.iter.as_legacy_range_mut();
486                let remaining = r.end;
487                if remaining > 0 {
488                    let val = r.start;
489                    // this can only overflow during the last step, after which the value
490                    // will not be used
491                    r.start = val.wrapping_add(step);
492                    r.end = remaining - 1;
493                    Some(val)
494                } else {
495                    None
496                }
497            }
498
499            #[inline]
500            fn spec_size_hint(&self) -> (usize, Option<usize>) {
501                let remaining = self.iter.as_legacy_range().end as usize;
502                (remaining, Some(remaining))
503            }
504
505            // The methods below are all copied from the Iterator trait default impls.
506            // We have to repeat them here so that the specialization overrides the StepByImpl defaults
507
508            #[inline]
509            fn spec_nth(&mut self, n: usize) -> Option<Self::Item> {
510                self.advance_by(n).ok()?;
511                self.next()
512            }
513
514            #[inline]
515            fn spec_try_fold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R
516                where
517                    F: FnMut(Acc, Self::Item) -> R,
518                    R: Try<Output = Acc>
519            {
520                let mut accum = init;
521                while let Some(x) = self.next() {
522                    accum = f(accum, x)?;
523                }
524                try { accum }
525            }
526
527            #[inline]
528            fn spec_fold<Acc, F>(self, init: Acc, mut f: F) -> Acc
529                where
530                    F: FnMut(Acc, Self::Item) -> Acc
531            {
532                // if a step size larger than the type has been specified fall back to
533                // t::MAX, in which case remaining will be at most 1.
534                let step = <$t>::try_from(self.original_step().get()).unwrap_or(<$t>::MAX);
535                let r = self.iter.as_legacy_range();
536                let remaining = r.end;
537                let mut acc = init;
538                let mut val = r.start;
539                for _ in 0..remaining {
540                    acc = f(acc, val);
541                    // this can only overflow during the last step, after which the value
542                    // will no longer be used
543                    val = val.wrapping_add(step);
544                }
545                acc
546            }
547        }
548    )*)
549}
550
551macro_rules! spec_int_ranges_r {
552    ($ctor:ident; $($t:ty)*) => ($(
553        const _: () = assert!(usize::BITS >= <$t>::BITS);
554
555        unsafe impl StepByBackImpl<$ctor<$t>> for StepBy<$ctor<$t>> {
556
557            #[inline]
558            fn spec_next_back(&mut self) -> Option<Self::Item> {
559                let step = self.original_step().get() as $t;
560                let r = self.iter.as_legacy_range_mut();
561                let remaining = r.end;
562                if remaining > 0 {
563                    let start = r.start;
564                    r.end = remaining - 1;
565                    Some(start + step * (remaining - 1))
566                } else {
567                    None
568                }
569            }
570
571            // The methods below are all copied from the Iterator trait default impls.
572            // We have to repeat them here so that the specialization overrides the StepByImplBack defaults
573
574            #[inline]
575            fn spec_nth_back(&mut self, n: usize) -> Option<Self::Item> {
576                if self.advance_back_by(n).is_err() {
577                    return None;
578                }
579                self.next_back()
580            }
581
582            #[inline]
583            fn spec_try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R
584            where
585                F: FnMut(Acc, Self::Item) -> R,
586                R: Try<Output = Acc>
587            {
588                let mut accum = init;
589                while let Some(x) = self.next_back() {
590                    accum = f(accum, x)?;
591                }
592                try { accum }
593            }
594
595            #[inline]
596            fn spec_rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc
597            where
598                F: FnMut(Acc, Self::Item) -> Acc
599            {
600                let mut accum = init;
601                while let Some(x) = self.next_back() {
602                    accum = f(accum, x);
603                }
604                accum
605            }
606        }
607    )*)
608}
609
610// The same specialization covers `Range<{integer}>` and the new-range iterator
611// `RangeIter<{integer}>`, which wraps a `Range` (see `AsLegacyRange`).
612//
613// The backward (`_r`) specialization requires `ExactSizeIterator`. `RangeIter`
614// implements it only for `usize`/`u8`/`u16` (see `range_exact_iter_impl!` in
615// `range::iter`), narrower than `Range`, so `RangeIter`'s backward set omits
616// `u32` even where `Range` includes it; `Range<u64>` is likewise omitted on
617// 64-bit since its length can exceed `usize`.
618#[cfg(target_pointer_width = "64")]
619mod step_by_spec {
620    use super::*;
621    spec_int_ranges!(Range; u8 u16 u32 u64 usize);
622    spec_int_ranges!(RangeIter; u8 u16 u32 u64 usize);
623    spec_int_ranges_r!(Range; u8 u16 u32 usize);
624    spec_int_ranges_r!(RangeIter; u8 u16 usize);
625}
626
627#[cfg(target_pointer_width = "32")]
628mod step_by_spec {
629    use super::*;
630    spec_int_ranges!(Range; u8 u16 u32 usize);
631    spec_int_ranges!(RangeIter; u8 u16 u32 usize);
632    spec_int_ranges_r!(Range; u8 u16 u32 usize);
633    spec_int_ranges_r!(RangeIter; u8 u16 usize);
634}
635
636#[cfg(target_pointer_width = "16")]
637mod step_by_spec {
638    use super::*;
639    spec_int_ranges!(Range; u8 u16 usize);
640    spec_int_ranges!(RangeIter; u8 u16 usize);
641    spec_int_ranges_r!(Range; u8 u16 usize);
642    spec_int_ranges_r!(RangeIter; u8 u16 usize);
643}
core/iter/adapters/step_by.rs

core/iter/adapters/
step_by.rs
